package com.wsw.aitools.test;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import org.springframework.http.*;
import org.springframework.util.Base64Utils;
import org.springframework.web.client.RestTemplate;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.Map;

public class DoubaoDemo {

    public static void main(String[] args) throws Exception{
        call();
    }
    
    private static void call() throws Exception{
        String prompt = "##角色：\n" +
                "你是一名中文图片分析专家，请你请按照以下规则分析并处理图片：\n" +
                "##任务\n" +
                "###任务1：如果图片中有字幕或台词，请从其中提取出3个以内的专业词汇；你需要优先选择其中的中文词汇；请注意，如果字幕过长，你只需要提取其中专业性最强的3个词语即可。\n" +
                "###任务2：根据图片内容推荐10部相关影片，以列表形式给出，并给出30到90个字的推荐理由。\n" +
                "##输出格式：\n" +
                "要求以json格式进行输出，格式如下：\n" +
                "示例：\n" +
                "{\"推荐理由\":\"推荐理由\",\"词语提取\": [ \"词语1\", \"词语2\", \"词语3\"], \"推荐\": [ \"影片1\", \"影片2\", \"影片3\",\"影片4\", ,\"影片5\",\"影片6\",\"影片7\",\"影片8\",\"影片9\",\"影片10\"]}\n" +
                "\n" +
                "##限制：\n" +
                "1.需要严格遵循json格式的要求\n" +
                "2.如果影片名称中可能涉及不文明词语，请不要使用该影片\n" +
                "3.请注意，你只需遵照格式给出相关内容，不要输出其他内容！";


//        String prompt = "请尽量清楚识别图片中包含的球员和球衣号码";
        
        String url = "https://ark.cn-beijing.volces.com/api/v3/chat/completions";
        RestTemplate restTemplate = new RestTemplate();

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth("158f035c-1d3e-4b0d-bf1c-63edef2fe1fd ");

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", "ep-20241224182005-8vh2m");

        Map<String, Object> textContent = new HashMap<>();
        textContent.put("type", "text");
        textContent.put("text", prompt);
        
//        String image_path = "img/老虎.png";
//        String image_path = "img/音乐综艺.png";
//        String image_path = "img/涉谷十字路口.png";
        String image_path = "img/熊猫.jpg";
//        String image_path = "img/足球2.jpg";
        // 对图片进行base64编码
        String s = encodeImage(image_path);


        Map<String, Object> imageUrl1 = new HashMap<>();
        imageUrl1.put("type", "image_url");
        imageUrl1.put("image_url", new HashMap<String, String>() {{
//            put("url", "https://ark-project.tos-cn-beijing.volces.com/images/view.jpeg");
            
            put("url", s);
        }});

//        Map<String, Object> imageUrl2 = new HashMap<>();
//        imageUrl2.put("type", "image_url");
//        imageUrl2.put("image_url", new HashMap<String, String>() {{
//            put("url", "https://portal.volccdn.com/obj/volcfe/cloud-universal-doc/upload_a81e3cdd3e30617ecd524a132fdb2736.png");
//        }});

        Map<String, Object> message = new HashMap<>();
        message.put("role", "user");
        message.put("content", new Object[]{textContent, imageUrl1});

        requestBody.put("messages", new Object[]{message});

        Long startTime = System.currentTimeMillis();
        HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);

        ResponseEntity<String> response = restTemplate.exchange(url, HttpMethod.POST, entity, String.class);
        
        System.out.println("Time: " + (System.currentTimeMillis() - startTime));
        
        String body = response.getBody();
        System.out.println("Response: " + response.getBody());
        
        // 转fastjson
        JSONObject jsonObject = JSON.parseObject(body);
        JSONArray choices = jsonObject.getJSONArray("choices");

        JSONObject choice1  = (JSONObject) choices.get(0);
        JSONObject messageJson = choice1.getJSONObject("message");
        String content = messageJson.getString("content");
        System.out.println("content: " + content);


    }

    public static String encodeImage(String imagePath) throws IOException {
        byte[] imageBytes = Files.readAllBytes(Paths.get(imagePath));
        String base64 = Base64Utils.encodeToString(imageBytes);
        
        // 需要注意：传入Base64编码前需要增加前缀 data:image/{图片格式};base64,{Base64编码}：
        // PNG图片："url":  f"data:image/png;base64,{base64_image}"
        // JEPG图片："url":  f"data:image/jpeg;base64,{base64_image}"
        // WEBP图片："url":  f"data:image/webp;base64,{base64_image}"
        
        // 判断图片格式
        String format = imagePath.substring(imagePath.lastIndexOf(".") + 1);
        switch (format) {
            case "png":
                return "data:image/png;base64," + base64;
            case "jpeg":
                return "data:image/jpeg;base64," + base64;
            case "jpg":
                return "data:image/jpg;base64," + base64;
            case "webp":
                return "data:image/webp;base64," + base64;
            default:
                return "data:image/png;base64," + base64;
        }
    }
}
